class DDoSDataset(Dataset):
    """DDoS攻击数据集类"""
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels
        
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]

def load_cicddos2019():
    """加载CICDDoS2019数据集（高频攻击）"""
    # 实际实现应从文件加载预处理数据
    # 这里使用模拟数据
    num_samples = 800000
    num_features = 83
    features = np.random.randn(num_samples, 1, num_features).astype(np.float32)
    labels = np.random.randint(0, 2, num_samples).astype(np.int64)
    return features, labels

def load_cicids2017():
    """加载CICIDS2017数据集（低频攻击）"""
    # 实际实现应从文件加载预处理数据
    # 这里使用模拟数据
    num_samples = 1200000
    num_features = 78
    features = np.random.randn(num_samples, 1, num_features).astype(np.float32)
    labels = np.random.randint(0, 2, num_samples).astype(np.int64)
    return features, labels

def create_non_iid_datasets(config):
    """创建Non-IID数据集"""
    # 加载数据集
    hr_features, hr_labels = load_cicddos2019()
    lr_features, lr_labels = load_cicids2017()
    
    # 合并数据集
    all_features = np.concatenate([hr_features, lr_features], axis=0)
    all_labels = np.concatenate([hr_labels, lr_labels], axis=0)
    
    # 创建客户端数据集
    client_datasets = []
    total_samples = len(all_labels)
    samples_per_client = total_samples // config.num_clients
    
    for i in range(config.num_clients):
        # 根据Non-IID程度创建有偏数据分布
        bias = config.non_iid_degree * i / config.num_clients
        start_idx = i * samples_per_client
        end_idx = (i + 1) * samples_per_client
        
        # 应用偏置
        biased_indices = np.arange(start_idx, end_idx)
        np.random.shuffle(biased_indices)
        bias_point = int(len(biased_indices) * bias)
        
        client_features = all_features[biased_indices[:bias_point]]
        client_labels = all_labels[biased_indices[:bias_point]]
        
        # 创建数据集
        dataset = DDoSDataset(client_features, client_labels)
        client_datasets.append(dataset)
    
    return client_datasets